# encoding=UTF-8
# https://blog.csdn.net/u011285477/article/details/52004513
import collections

import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt


def Otsu(gray_img: np.ndarray, REG: float = 1e-1):
    """
        大津法OTSU算法：

            算法思想：

                根据直方图均衡化后的概率直方图（一维数组）
                从数组左边开始不断累加，并调用方差计算公式
                找到让方差最大的像素值(阈值)并返回

            算法适用范围：

                适用于双峰直方图，原理就是算法思想中就是分割背景+前景而提出的
                至于双峰，可以统计图像的直方图，看看是否大致符合双峰(直接matplotlib画出来)

    :param gray_img: 虽然有防错机制，但是务必传入灰度图像
    :return: threshold, variance(可以用于前置算法调参)
    """
    ex = 0
    if gray_img.size == 0:
        ex = Exception("图像数组为空，请检查图像输入")
    elif len(gray_img.shape) != 2:
        ex = Exception("图像非灰度图像")
    if ex != 0:
        raise ex

    if REG == 0:
        REG = 1e-3

    # 实际上灰度图的channel=1,写着就当防止报错了
    # 但是传入彩色图，会对精度产生影响
    width, height = gray_img.shape

    # 统计图像直方图
    value_count = collections.Counter(gray_img.reshape((width * height)))
    histogram = np.array([0] * 256)
    for item in value_count.items():
        index, count = item
        histogram[index] = count

    # normalize
    histogram = histogram / (width * height)

    # 给每个位置添加权重并统计平均灰度(u)
    # 实际上就是期望，weight代表[0-255]，histogram代表每个值出现的概率
    # 所以这里的命名有失误的地方
    weight = np.array([i for i in range(0, 256)])
    u = sum(histogram * weight)

    # 计算阈值
    # 算法遍历0-255的颜色范围
    # 假定前景色占比(w0:权重)为循环的累加值
    # 利用公式：(u0 - u)^2 * w0 / (1-w0)
    threshold = 0
    max_variance = 0
    w0 = 0
    avg_value = 0
    for i in range(0, 256):
        w0 += histogram[i]  # 前景的权重
        avg_value += i * histogram[i]  # avg_value = w0 * u0, 前景的概率期望

        t = (avg_value + REG) / (w0 + REG) - u  # u0 - u, 1e-12是为了防止0/0,并且自身足够小，不影响精度
        temp_variance = (t * t * w0 + REG) / (1 - w0 + REG)
        if temp_variance > max_variance:
            max_variance = temp_variance
            threshold = i

    return threshold, max_variance


if __name__ == '__main__':
    # 文件读取与转换
    img = cv.imread("../resources/Lenna.png")
    gray = cv.cvtColor(img, cv.COLOR_RGB2GRAY)

    # 计算
    th, var = Otsu(gray)

    # 二值化
    retval, dst = cv.threshold(gray, th, 255, cv.THRESH_BINARY)

    cv.imshow("es", dst)
    cv.waitKey(0)